`include "all_define.vh"

module mem_stage (
  input clk,
  input rst_n,

  input                        exe_mem_vld,
  input                        exe_mem_opena,
  input                        exe_mem_rwn,
  input [`MEM_DTYPE_WIDTH-1:0] exe_mem_dtype,
  input [31:0]                 exe_mem_addr,
  input [31:0]                 exe_mem_datain,

  input                        exe_mem_wb_ena,
  input [4:0]                  exe_mem_wb_rd,
  input [31:0]                 exe_mem_wb_rd_data,

  output                       dmem_if_req,
  output                       dmem_if_rwn,
  output     [1:0]             dmem_if_size,
  output     [31:0]            dmem_if_addr,
  output     [31:0]            dmem_if_wdata,
  input                        dmem_if_grant,
  input      [31:0]            dmem_if_rdata,

  output reg                   mem_wb_vld,
  output reg                   mem_wb_is_opload,
  output reg [4:0]             mem_wb_rd,
  output reg [31:0]            mem_wb_rd_data,
  output reg [31:0]            mem_wb_mrdata,

  output                       mem_stall_req,

  output [4:0]                 mem_bypass_rd,
  output                       mem_bypass_rd_ena,
  output [31:0]                mem_bypass_rd_data,
  output [31:0]                mem_bypass_load_data,
  output                       mem_bypass_load_data_vld
);

  wire   mem_op_load;

  assign mem_op_load = exe_mem_vld & exe_mem_opena & exe_mem_rwn;

  assign dmem_if_req   = exe_mem_vld & exe_mem_opena;
  assign dmem_if_rwn   = exe_mem_rwn;
  assign dmem_if_size  = exe_mem_dtype[1:0];
  assign dmem_if_addr  = exe_mem_addr;
  assign dmem_if_wdata = exe_mem_datain;

  wire [7:0] rdata_byte;
  wire [15:0] rdata_half;
  assign rdata_byte = dmem_if_rdata[7:0];
  assign rdata_half = dmem_if_rdata[15:0];

  reg [31:0] rdata;
  always @(*) begin
    case (exe_mem_dtype[1:0])
      2'b00: rdata = exe_mem_dtype[2]? {24'b0, rdata_byte}: {{24{rdata_byte[7]}}, rdata_byte};
      2'b01: rdata = exe_mem_dtype[2]? {16'b0, rdata_half}: {{16{rdata_half[15]}}, rdata_half};
      2'b10: rdata = dmem_if_rdata;
      default: rdata = dmem_if_rdata;
    endcase
  end

  always @(*) begin
    mem_wb_vld = 1'b0;
    if (exe_mem_vld) begin
      if (mem_op_load) begin
        mem_wb_vld = dmem_if_grant & exe_mem_wb_ena;
      end else begin
        mem_wb_vld = exe_mem_wb_ena;
      end
    end
  end

  always @(*) begin
    mem_wb_rd = exe_mem_wb_rd;
    mem_wb_rd_data = exe_mem_wb_rd_data;
    mem_wb_mrdata = rdata;
    mem_wb_is_opload = mem_op_load;
  end

  // request for stall pipeline if memory operation doesn't complete
  assign mem_stall_req = dmem_if_req & (~dmem_if_grant);

  assign mem_bypass_rd      = exe_mem_wb_rd;
  assign mem_bypass_rd_ena  = (~mem_op_load)& exe_mem_vld & exe_mem_wb_ena;
  assign mem_bypass_rd_data = exe_mem_wb_rd_data;
  assign mem_bypass_load_data     = rdata;
  assign mem_bypass_load_data_vld = mem_op_load & dmem_if_grant;

endmodule

